import json
from io import BytesIO
import numpy as np
from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import base64
import os
import analyse


def kmeans_cluster():
    # 打开JSON文件
    model_dir = os.path.dirname(analyse.__file__)
    data_path = os.path.join(model_dir, 'data', 'job_train_data.json')
    with open(data_path, 'r', encoding='utf-8') as json_file:
        data = json.load(json_file)

    # 选择两个特征，例如城市评分和薪水上限
    feature1 = 'city_score'
    feature2 = 'salary_upper_limit'

    # 提取所选特征的值
    X = np.array([[item[feature1], item[feature2]] for item in data])

    # 执行K均值聚类
    kmeans = KMeans(n_clusters=3, n_init='auto')  # 选择聚类数量
    kmeans.fit(X)

    # 使用t-SNE进行降维
    tsne = TSNE(n_components=2)  # 选择目标维度为2
    X_tsne = tsne.fit_transform(X)

    # 绘制聚类图
    plt.figure(figsize=(10, 6))
    for cluster_label in set(kmeans.labels_):
        cluster_data = X_tsne[kmeans.labels_ == cluster_label]
        plt.scatter(cluster_data[:, 0], cluster_data[:, 1], label=f'Cluster {cluster_label}', s=50)

    plt.xlabel('t-SNE city_code')
    plt.ylabel('t-SNE salary_lower_limit')
    plt.title('K-Means')
    plt.legend()

    # 生成图像的Base64编码
    img_data = BytesIO()
    plt.savefig(img_data, format='png')
    img_data.seek(0)
    img_base64 = base64.b64encode(img_data.read()).decode('utf-8')

    # 关闭图形
    plt.close()
    return img_base64
